#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* 包含动态库头文件 */
#include "libtwclml.h"

/**
 * 打印设备信息
 */
void get_device_info(nvmlDevice_t device) {
    if (!device) return;
    
    unsigned int index;
    nvmlPciInfo_t pci;
    int major, minor;
    
    nvmlDeviceGetIndex(device, &index);
    nvmlDeviceGetPciInfo_v2(device, &pci);
    nvmlDeviceGetCudaComputeCapability(device, &major, &minor);
    
    printf("设备索引: %u\n", index);
    printf("PCI信息: %s (domain: 0x%04X, bus: 0x%02X, device: 0x%02X, function: 0x%01X)\n",
           pci.busId, pci.domain, pci.bus, pci.device, pci.function);
    printf("计算能力: %d.%d\n", major, minor);
}

/**
 * 测试设备计数功能
 */
int test_device_count() {
    unsigned int count;
    nvmlReturn_t ret = nvmlDeviceGetCount(&count);
    
    if (ret != NVML_SUCCESS) {
        printf("获取设备数量失败: %s\n", nvmlErrorString(ret));
        return -1;
    }
    
    printf("检测到 %u 个设备\n", count);
    return count;
}

/**
 * 测试设备信息获取功能
 */
int test_device_info(unsigned int index) {
    nvmlDevice_t device;
    nvmlReturn_t ret = nvmlDeviceGetHandleByIndex(index, &device);
    
    if (ret != NVML_SUCCESS) {
        printf("获取设备 %u 句柄失败: %s\n", index, nvmlErrorString(ret));
        return -1;
    }
    
    printf("\n----- 设备 %u 信息 ----- \n", index);
    get_device_info(device);
    
    // 测试通过PCI Bus ID获取设备句柄
    nvmlPciInfo_t pci;
    nvmlDeviceGetPciInfo_v2(device, &pci);
    
    nvmlDevice_t device_by_bus;
    ret = nvmlDeviceGetHandleByPciBusId(pci.busId, &device_by_bus);
    if (ret != NVML_SUCCESS) {
        printf("通过PCI Bus ID获取设备句柄失败: %s\n", nvmlErrorString(ret));
    } else {
        unsigned int idx;
        nvmlDeviceGetIndex(device_by_bus, &idx);
        printf("通过PCI Bus ID %s 成功获取设备句柄，索引: %u\n", pci.busId, idx);
    }
    
    return 0;
}

/**
 * 测试NVLink信息获取功能
 */
int test_nvlink_info(nvmlDevice_t device) {
    if (!device) return -1;
    
    unsigned int index;
    unsigned int capability;
    nvmlDeviceGetIndex(device, &index);
    
    // 获取NVLink数量（简化实现，实际应从设备属性获取）
    unsigned int num_links = 4; // 假设最多4个NVLink
    
    printf("\n----- 设备 %u 的NVLink信息 ----- \n", index);
    
    for (unsigned int link = 0; link < num_links; link++) {
        nvmlEnableState_t state;
        nvmlReturn_t ret = nvmlDeviceGetNvLinkState(device, link, &state);
        
        if (ret != NVML_SUCCESS) {
            printf("获取设备 %u 的NVLink %u 状态失败: %s\n", 
                   index, link, nvmlErrorString(ret));
            continue;
        }
        
        printf("NVLink %u: %s\n", link, state == NVML_ENABLED ? "激活" : "未激活");
        
        if (state == NVML_ENABLED) {
            nvmlPciInfo_t remote_pci;
            ret = nvmlDeviceGetNvLinkRemotePciInfo(device, link, &remote_pci);
            
            if (ret == NVML_SUCCESS) {
                printf("  远程设备PCI信息: %s\n", remote_pci.busId);
            }
        }
    // 检查P2P支持能力
        ret = nvmlDeviceGetNvLinkCapability(device, link,
                                            NVML_NVLINK_CAP_P2P_SUPPORTED, &capability);
        if (ret != 0) {
            printf("  链路 %u: 获取P2P支持能力失败: %s\n", link, nvmlErrorString(ret));
        } else {
            printf("  链路 %u: P2P支持=%s\n", link,
                    capability ? "是" : "否");
        }

    }
    
    return 0;
}

/**
 * 测试P2P能力获取功能
 */
int test_p2p_info(nvmlDevice_t device1, nvmlDevice_t device2) {
    if (!device1 || !device2) return -1;
    
    unsigned int index1, index2;
    nvmlDeviceGetIndex(device1, &index1);
    nvmlDeviceGetIndex(device2, &index2);
    
    if (index1 == index2) return 0; // 跳过与自身的测试
    
    printf("\n----- 设备 %u 与设备 %u 的P2P能力 ----- \n", index1, index2);
    
    nvmlGpuP2PStatus_t status;
    
    // 测试读能力
    nvmlReturn_t ret = nvmlDeviceGetP2PStatus(device1, device2, 
                                            NVML_P2P_CAPS_INDEX_READ, &status);
    if (ret == NVML_SUCCESS) {
        printf("读操作支持: %s\n", status == NVML_P2P_STATUS_OK ? "是" : "否");
    }
    
    // 测试写能力
    ret = nvmlDeviceGetP2PStatus(device1, device2, 
                               NVML_P2P_CAPS_INDEX_WRITE, &status);
    if (ret == NVML_SUCCESS) {
        printf("写操作支持: %s\n", status == NVML_P2P_STATUS_OK ? "是" : "否");
    }
    
    // 测试NVLink支持
    ret = nvmlDeviceGetP2PStatus(device1, device2, 
                               NVML_P2P_CAPS_INDEX_NVLINK, &status);
    if (ret == NVML_SUCCESS) {
        printf("NVLink支持: %s\n", status == NVML_P2P_STATUS_OK ? "是" : "否");
    }
    
    // 测试原子操作支持
    ret = nvmlDeviceGetP2PStatus(device1, device2, 
                               NVML_P2P_CAPS_INDEX_ATOMICS, &status);
    if (ret == NVML_SUCCESS) {
        printf("原子操作支持: %s\n", status == NVML_P2P_STATUS_OK ? "是" : "否");
    }
    
    return 0;
}

int main() {
    nvmlReturn_t ret;
    unsigned int device_count;
    nvmlDevice_t* devices;
    
    /* 初始化库 */
    ret = nvmlInit();
    if (ret != NVML_SUCCESS) {
        printf("初始化库失败: %s\n", nvmlErrorString(ret));
        return EXIT_FAILURE;
    }
    printf("成功初始化TWCLML库\n");
    
    /* 测试设备计数 */
    device_count = test_device_count();
    if (device_count <= 0) {
        nvmlShutdown();
        return EXIT_FAILURE;
    }
    
    /* 分配设备句柄数组 */
    devices = (nvmlDevice_t*)malloc(sizeof(nvmlDevice_t) * device_count);
    if (!devices) {
        printf("内存分配失败\n");
        nvmlShutdown();
        return EXIT_FAILURE;
    }
    
    /* 获取所有设备句柄并测试设备信息 */
    printf("\n===== 设备基本信息测试 =====");
    for (unsigned int i = 0; i < device_count; i++) {
        ret = nvmlDeviceGetHandleByIndex(i, &devices[i]);
        if (ret != NVML_SUCCESS) {
            printf("获取设备 %u 句柄失败: %s\n", i, nvmlErrorString(ret));
            continue;
        }
        test_device_info(i);
    }
    
    /* 测试NVLink信息 */
    printf("\n===== NVLink 信息测试 =====");
    for (unsigned int i = 0; i < device_count; i++) {
        if (devices[i]) {
            test_nvlink_info(devices[i]);
        }
    }
    
    /* 测试P2P能力 */
    printf("\n===== P2P 能力测试 =====");
    for (unsigned int i = 0; i < device_count; i++) {
        for (unsigned int j = 0; j < device_count; j++) {
            if (devices[i] && devices[j]) {
                test_p2p_info(devices[i], devices[j]);
            }
        }
    }
    
    /* 释放资源 */
    for (unsigned int i = 0; i < device_count; i++) {
        if (devices[i]) {
            free(devices[i]);
        }
    }
    free(devices);
    
    /* 关闭库 */
    ret = nvmlShutdown();
    if (ret != NVML_SUCCESS) {
        printf("关闭库失败: %s\n", nvmlErrorString(ret));
        return EXIT_FAILURE;
    }
    printf("\n成功关闭TWCLML库\n");
    
    return EXIT_SUCCESS;
}
